import math
import logging
import oneflow.nn as nn


logger = logging.getLogger(__name__)


def initiailze_conv_params(n, p):
    """Initialize with Lecun style.

    Args:
        n (str): parameter name
        p (Tensor): parameter
    """
    if p.dim() == 1:
        nn.init.constant_(p, 0.)  # bias
    elif p.dim() == 2:
        fan_in = p.size(1)
        nn.init.normal_(p, mean=0., std=1. / math.sqrt(fan_in))  # linear weight
    elif p.dim() == 3:
        fan_in = p.size(1) * p[0][0].numel()
        nn.init.normal_(p, mean=0., std=1. / math.sqrt(fan_in))  # 1d conv weight
    elif p.dim() == 4:
        fan_in = p.size(1) * p[0][0].numel()
        nn.init.normal_(p, mean=0., std=1. / math.sqrt(fan_in))  # 2d conv weight
    else:
        raise ValueError(n)


def initiailze_linear_params(n, p):
    if p.dim() == 1:
        nn.init.constant_(p, 0.)  # bias
    elif p.dim() in [2, 3, 4]:
        nn.init.xavier_uniform_(p)
    else:
        raise ValueError(n)
